package com.dms.common.security;

import com.dms.common.config.JwtConfig;
import io.jsonwebtoken.*;
import io.jsonwebtoken.security.Keys;
import io.jsonwebtoken.security.SignatureException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.security.core.userdetails.UserDetailsService;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.servlet.http.HttpServletRequest;
import java.util.Base64;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.Arrays;

@Slf4j
@Component
public class JwtTokenProvider {

    @Autowired
    private JwtConfig jwtConfig;

    @Autowired
    @Qualifier("userServiceImpl")
    private UserDetailsService userDetailsService;

    @Autowired
    @Qualifier("merchantUserDetailsService")
    private UserDetailsService merchantUserDetailsService;

    private String secretKey;

    @PostConstruct
    protected void init() {
        secretKey = Base64.getEncoder().encodeToString(jwtConfig.getSecret().getBytes());
    }

    /**
     * 生成标准JWT令牌
     */
    public String generateToken(Authentication authentication) {
        UserDetails userDetails = (UserDetails) authentication.getPrincipal();
        return generateToken(userDetails);
    }

    /**
     * 从UserDetails生成JWT令牌
     */
    public String generateToken(UserDetails userDetails) {
        Map<String, Object> claims = new HashMap<>();
        claims.put("roles", userDetails.getAuthorities().stream()
                .map(GrantedAuthority::getAuthority)
                .collect(Collectors.toList()));
                
        Date now = new Date();
        Date validity = new Date(now.getTime() + jwtConfig.getAccessTokenExpiration());
        
        return Jwts.builder()
                .setClaims(claims)
                .setSubject(userDetails.getUsername())
                .setIssuedAt(now)
                .setExpiration(validity)
                .signWith(Keys.hmacShaKeyFor(secretKey.getBytes()), SignatureAlgorithm.HS256)
                .compact();
    }

    /**
     * 生成商家专用令牌
     */
    public String generateMerchantToken(Long staffId, Long merchantId) {
        Map<String, Object> claims = new HashMap<>();
        claims.put("staffId", staffId);
        claims.put("merchantId", merchantId);
        claims.put("roles", Arrays.asList("ROLE_MERCHANT"));
        
        Date now = new Date();
        Date validity = new Date(now.getTime() + jwtConfig.getAccessTokenExpiration());
        
        return Jwts.builder()
                .setClaims(claims)
                .setSubject(staffId.toString())
                .setIssuedAt(now)
                .setExpiration(validity)
                .signWith(Keys.hmacShaKeyFor(secretKey.getBytes()), SignatureAlgorithm.HS256)
                .compact();
    }
    
    /**
     * 生成刷新令牌
     */
    public String generateRefreshToken(Authentication authentication) {
        UserDetails userDetails = (UserDetails) authentication.getPrincipal();
        Map<String, Object> claims = new HashMap<>();
        claims.put("type", "refresh");
        claims.put("roles", userDetails.getAuthorities().stream()
                .map(GrantedAuthority::getAuthority)
                .collect(Collectors.toList()));

        Date now = new Date();
        Date validity = new Date(now.getTime() + jwtConfig.getRefreshTokenExpiration());

        return Jwts.builder()
                .setClaims(claims)
                .setSubject(userDetails.getUsername())
                .setIssuedAt(now)
                .setExpiration(validity)
                .signWith(Keys.hmacShaKeyFor(secretKey.getBytes()), SignatureAlgorithm.HS256)
                .compact();
    }

    public Authentication getAuthentication(String token) {
        Claims claims = getAllClaims(token);
        
        // 获取用户名和角色
        String username = claims.getSubject();
        List<String> roles = claims.get("roles", List.class);
        
        if (username == null || roles == null) {
            log.warn("Invalid JWT token: missing username or roles");
            return null;
        }

        // 如果是商家令牌，需要特殊处理
        Long merchantId = claims.get("merchantId", Long.class);
        if (merchantId != null) {
            try {
                UserDetails userDetails = merchantUserDetailsService.loadUserByUsername(username);
                return new UsernamePasswordAuthenticationToken(userDetails, "", userDetails.getAuthorities());
            } catch (Exception e) {
                log.error("Failed to load merchant: {}", username, e);
                return null;
            }
        }
        
        // 普通用户认证
        try {
            UserDetails userDetails = userDetailsService.loadUserByUsername(username);
            return new UsernamePasswordAuthenticationToken(userDetails, "", userDetails.getAuthorities());
        } catch (Exception e) {
            log.error("Failed to load user: {}", username, e);
            return null;
        }
    }

    public String getUsernameFromToken(String token) {
        return extractClaim(token, Claims::getSubject);
    }

    public Long getMerchantIdFromToken(String token) {
        Claims claims = getAllClaims(token);
        return claims.get("merchantId", Long.class);
    }

    public Long getStaffIdFromToken(String token) {
        Claims claims = getAllClaims(token);
        return claims.get("staffId", Long.class);
    }
    
    public <T> T extractClaim(String token, Function<Claims, T> claimsResolver) {
        final Claims claims = getAllClaims(token);
        return claimsResolver.apply(claims);
    }
    
    private Claims getAllClaims(String token) {
        return Jwts.parserBuilder()
                .setSigningKey(Keys.hmacShaKeyFor(secretKey.getBytes()))
                .build()
                .parseClaimsJws(token)
                .getBody();
    }
    
    public String extractUsername(String token) {
        return extractClaim(token, Claims::getSubject);
    }
    
    public Boolean isRefreshToken(String token) {
        try {
            Claims claims = getAllClaims(token);
            return "refresh".equals(claims.get("type"));
        } catch (Exception e) {
            return false;
        }
    }

    public boolean validateToken(String token) {
        try {
            Jwts.parserBuilder()
                .setSigningKey(Keys.hmacShaKeyFor(secretKey.getBytes()))
                .build()
                .parseClaimsJws(token);
            return true;
        } catch (JwtException | IllegalArgumentException e) {
            log.error("Invalid JWT token: {}", e.getMessage());
            return false;
        }
    }

    public String resolveToken(HttpServletRequest request) {
        String bearerToken = request.getHeader("Authorization");
        if (bearerToken != null && bearerToken.startsWith("Bearer ")) {
            return bearerToken.substring(7);
        }
        return null;
    }

    public Long getAccessTokenExpiration() {
        return jwtConfig.getAccessTokenExpiration();
    }
    
    public Long getRefreshTokenExpiration() {
        return jwtConfig.getRefreshTokenExpiration();
    }
} 